{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "321adc7f-f386-4dbf-bb9c-d5f789421f78",
   "metadata": {},
   "source": [
    "# Build and Deploy an Agent Development Kit (ADK) Agent\n",
    "\n",
    "> IMPORTANT: This is the third notebook in the lab. The notebooks  build on top of each other, so be sure to run the preceding notebooks, in order, before running this one. Start your journey building ADK Agents with MCP Toolbox [here](./1_setup_and_explore_databases.ipynb). \n",
    "\n",
    "## Overview\n",
    "\n",
    "<img src=\"img/agent_development_kit_logo.png\" align=\"left\" style=\"height: 140px;\"> Agent Development Kit (ADK) is a flexible and modular framework for developing and deploying AI agents. While optimized for Gemini and the Google ecosystem, ADK is model-agnostic, deployment-agnostic, and is built for compatibility with other frameworks. ADK was designed to make agent development feel more like software development, to make it easier for developers to create, deploy, and orchestrate agentic architectures that range from simple tasks to complex workflows.\n",
    "\n",
    "In this notebook, you will build and deploy a sophisticated, enterprise-ready AI agent using the Agent Development Kit (ADK) and MCP Toolbox for Databases. You will start by creating and testing an agent that runs entirely on the local Vertex AI Workbench instance, and then transition it to a scalable, enterprise-ready architecture by deploying it as a secure microservice in Cloud Run. \n",
    "\n",
    "The key milestones you will achieve are:\n",
    "\n",
    "- Building an Agent with MCP Toolbox: You will define and deploy an ADK agent that utilizes the MCP Toolbox to securely connect to the AlloyDB and Spanner databases provisioned in the previous labs. This will demonstrate how to create agents that can perform complex financial tasks, such as fraud detection, transaction lookups, and account analysis.\n",
    "- Secure Deployment to Cloud Run: You will learn how to deploy the ADK agent to a secure Cloud Run environment within a private VPC. The deployment will emphasize enterprise security by using a dedicated service account, managing secrets with Secret Manager, and enabling Cloud Logging and Cloud Trace for robust observability.\n",
    "- Enabling Telemetry and Observability: The notebook will guide you through enabling built-in tracing with ADK, allowing you to visualize the end-to-end execution of your agent's interactions, including tool calls made through the MCP Toolbox.\n",
    "- Interacting with a Deployed Agent: You will learn how to programmatically invoke the deployed agent.\n",
    "\n",
    "### Terraform Resources\n",
    "\n",
    "The following pre-requisite resources were created for you by Terraform. See the [main.tf](../terraform/main.tf) file for more details on the environment configuration:\n",
    "\n",
    "- Custom VPC (demo-vpc): Creates the demo-vpc private network where the Cloud Run services for the ADK agent and MCP Toolbox are deployed, ensuring secure communication between them.\n",
    "- Cloud Router and Cloud NAT: Enables the Cloud Run services to access external resources securely without exposing them to the public internet.\n",
    "- IAM Service Account: Defines the toolbox-service-account service account, which has the necessary permissions to access the AlloyDB and Spanner databases, as well as Secret Manager.\n",
    "- Cloud SQL for PostgreSQL Instance: The Cloud SQL instance that is used by the ADK agent for session management, ensuring that conversation states are persisted securely.\n",
    "- Cloud DNS: Sets up a private DNS zone to enable secure service-to-service communication between the ADK agent and MCP Toolbox within the VPC.\n",
    "\n",
    "### Google Cloud Services Used\n",
    "\n",
    "This notebook makes use of the following Google Cloud services:\n",
    "\n",
    "- Cloud Run: Hosts the ADK agent and MCP Toolbox as secure, scalable services within a private VPC.\n",
    "- Secret Manager: Securely stores the tools.yaml configuration and database credentials, which are accessed by the Cloud Run services at runtime.\n",
    "- IAM (Identity and Access Management): Manages permissions for the service accounts used by the ADK agent and MCP Toolbox, ensuring least-privilege access to resources.\n",
    "- Cloud SQL for PostgreSQL: Provides a persistent, managed database for the ADK agent's session management, enabling stateful conversations.\n",
    "- AlloyDB for PostgreSQL: Serves as a high-performance data source for the MCP Toolbox, allowing the agent to query financial transaction data.\n",
    "- Spanner: Acts as another critical data source, providing the agent with access to relational and graph data for complex queries.\n",
    "- MCP Toolbox for Databases (Toolbox): An open-source MCP (Model Context Protocol) server that allows developers to connect gen AI agents to enterprise data easily and securely.\n",
    "- Vertex AI Workbench: The development environment for running this notebook and interacting with the deployed services.\n",
    "- Cloud Logging and Cloud Trace: Provide observability into the agent's operations, allowing you to monitor logs and trace requests as they flow through the system.\n",
    "\n",
    "### Logical Flow\n",
    "\n",
    "The notebook is structured to guide you through the complete lifecycle of building and deploying an ADK agent with MCP Toolbox tools:\n",
    "- Basic Setup: This section prepares your environment by defining variables, connecting to your Google Cloud project, and installing the necessary Python libraries.\n",
    "- Build an Agent with MCP Toolbox Tools: Here, you will define the ADK agent, specifying its model, name, description, and instructions. Crucially, you will load a toolset from the deployed MCP Toolbox, enabling the agent to interact with your databases.\n",
    "- Manually Invoke the Agent: Before deploying, you will test the agent locally by sending it a series of prompts. This allows you to verify that the agent can correctly interpret the prompts and utilize the MCP Toolbox tools to retrieve information from the databases.\n",
    "- Inspect Agent Processing: To gain deeper insight into the agent's behavior, you will enable detailed logging. This will reveal the step-by-step execution of the agent, including the specific tool calls it makes, which is invaluable for debugging and understanding the agent's logic.\n",
    "- Deploy the ADK Agent to Cloud Run: This is a critical step where you will deploy the agent as a secure Cloud Run service. You will learn how to configure the service to use Cloud SQL for session persistence, ensuring that conversations are not lost upon server restarts. You will also enable Cloud Logging and Cloud Trace for observability.\n",
    "- Invoke the Cloud Run Agent: Once deployed, you will interact with the agent running on Cloud Run. The notebook demonstrates how to handle authentication and session management when invoking a deployed agent, providing a template for building interactive applications.\n",
    "- View Persisted ADK Data in Cloud SQL: To confirm that the session data is being persisted correctly, you will connect to the Cloud SQL instance and query the tables used by the ADK for session, event, and state management.\n",
    "- View Telemetry: Finally, you will explore the traces generated by the agent's invocations in Cloud Trace. This will provide a visual representation of the request flow, from the user's prompt to the agent's response, including the tool calls made to the MCP Toolbox."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0927ae02-125e-4d3e-a885-1e437e6bddb6",
   "metadata": {},
   "source": [
    "## Basic Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "842a3ca3-4f62-4190-bbb6-10e7da9b13c8",
   "metadata": {},
   "source": [
    "### Define Notebook Variables\n",
    "\n",
    "Update the `project_id` and `region` variables below to match your environment. You can use defaults for the rest of the project variables. \n",
    "\n",
    "You will be prompted for the Cloud SQL password you chose then you provisioned the environment with Terraform. We'll use the Cloud SQL instance in this notebook to store ADK Session state. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4fa431b-2efd-4db1-89e1-b97b4fc722df",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Project variables\n",
    "project_id = \"your-project\"\n",
    "region = \"your-region\"\n",
    "vpc = \"demo-vpc\"\n",
    "gcs_bucket_name = f\"project-files-{project_id}\"\n",
    "\n",
    "# Cloud SQL Variables\n",
    "cloud_sql_instance = \"my-postgres-instance\"\n",
    "cloud_sql_database = \"postgres\"\n",
    "cloud_sql_user = \"postgres\"\n",
    "cloud_sql_password = input(\"Please enter the password for the Cloud SQL 'postgres' database user: \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01db8a8b-cae4-4978-bb9a-e2e5a8465696",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set env variable to suppress annoying system warnings when running shell commands\n",
    "%env GRPC_ENABLE_FORK_SUPPORT=1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62bc39dc-f6ee-455e-853a-643c12918c2f",
   "metadata": {},
   "source": [
    "### Connect to your Google Cloud Project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d972baa-ffa7-47c7-a965-fdaa114262b5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Configure gcloud.\n",
    "!gcloud config set project {project_id}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38a124be-854e-4465-bab2-44ea855d9256",
   "metadata": {},
   "source": [
    "### Configure Logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66056e25-5e3a-47e3-a842-8fe6fadd9a33",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "# Configure the root logger to output messages with ERROR level\n",
    "logging.basicConfig(level=logging.ERROR, stream=sys.stdout, format='%(asctime)s[%(levelname)5s][%(name)14s] - %(message)s',  datefmt='%H:%M:%S', force=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b72eb606-f634-424f-be79-c1b75b1b1e14",
   "metadata": {},
   "source": [
    "### Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7c2e387-9dfc-49c2-8a32-2c54d7eebcc2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! pip install --quiet toolbox-core==0.2.1 \\\n",
    "                      google-adk==1.5.0 \\\n",
    "                      google-genai==1.23.0 \\\n",
    "                      asyncpg==0.30.0 \\\n",
    "                      \"cloud-sql-python-connector[asyncpg]\"==1.18.2 \\\n",
    "                      opentelemetry-instrumentation-fastapi \\\n",
    "                      opentelemetry-exporter-gcp-trace\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "052833d7-ccc3-470f-8974-7a741737eab0",
   "metadata": {},
   "source": [
    "### Define Helper Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5480756-e5fc-45eb-94c1-053ac52ce1d8",
   "metadata": {},
   "source": [
    "#### rest_api_helper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d90c79ab-d51f-42b2-bdb5-a6cbb469ef7e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "import google.auth\n",
    "import json\n",
    "\n",
    "# Get an access token based upon the current user\n",
    "creds, _ = google.auth.default()\n",
    "authed_session = google.auth.transport.requests.AuthorizedSession(creds)\n",
    "access_token=creds.token\n",
    "\n",
    "if project_id:\n",
    "  authed_session.headers.update({\"x-goog-user-project\": project_id}) # Required to workaround a project quota bug\n",
    "\n",
    "def rest_api_helper(\n",
    "    url: str,\n",
    "    http_verb: str,\n",
    "    request_body: dict = None,\n",
    "    params: dict = None,\n",
    "    session: requests.Session = authed_session,\n",
    "  ) -> dict:\n",
    "  \"\"\"Calls a REST API using a pre-authenticated requests Session.\"\"\"\n",
    "\n",
    "  headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "  try:\n",
    "\n",
    "    if http_verb == \"GET\":\n",
    "      response = session.get(url, headers=headers, params=params)\n",
    "    elif http_verb == \"POST\":\n",
    "      response = session.post(url, json=request_body, headers=headers, params=params)\n",
    "    elif http_verb == \"PUT\":\n",
    "      response = session.put(url, json=request_body, headers=headers, params=params)\n",
    "    elif http_verb == \"PATCH\":\n",
    "      response = session.patch(url, json=request_body, headers=headers, params=params)\n",
    "    elif http_verb == \"DELETE\":\n",
    "      response = session.delete(url, headers=headers, params=params)\n",
    "    else:\n",
    "      raise ValueError(f\"Unknown HTTP verb: {http_verb}\")\n",
    "\n",
    "    # Raise an exception for bad status codes (4xx or 5xx)\n",
    "    response.raise_for_status()\n",
    "\n",
    "    # Check if response has content before trying to parse JSON\n",
    "    if response.content:\n",
    "        return response.json()\n",
    "    else:\n",
    "        return {} # Return empty dict for empty responses (like 204 No Content)\n",
    "\n",
    "  except requests.exceptions.RequestException as e:\n",
    "      # Catch potential requests library errors (network, timeout, etc.)\n",
    "      # Log detailed error information\n",
    "      print(f\"Request failed: {e}\")\n",
    "      if e.response is not None:\n",
    "          print(f\"Request URL: {e.request.url}\")\n",
    "          print(f\"Request Headers: {e.request.headers}\")\n",
    "          print(f\"Request Body: {e.request.body}\")\n",
    "          print(f\"Response Status: {e.response.status_code}\")\n",
    "          print(f\"Response Text: {e.response.text}\")\n",
    "          # Re-raise a more specific error or a custom one\n",
    "          raise RuntimeError(f\"API call failed with status {e.response.status_code}: {e.response.text}\") from e\n",
    "      else:\n",
    "          raise RuntimeError(f\"API call failed: {e}\") from e\n",
    "  except json.JSONDecodeError as e:\n",
    "      print(f\"Failed to decode JSON response: {e}\")\n",
    "      print(f\"Response Text: {response.text}\")\n",
    "      raise RuntimeError(f\"Invalid JSON received from API: {response.text}\") from e\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "056b886f-c7b5-4699-9207-d73458b38445",
   "metadata": {},
   "source": [
    "#### get_auth_token()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ab17308-f483-4aff-95f9-33f3c910d652",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import urllib\n",
    "\n",
    "import google.auth.transport.requests\n",
    "import google.oauth2.id_token\n",
    "\n",
    "\n",
    "def get_auth_token(endpoint):\n",
    "    # Cloud Run uses your service's hostname as the `audience` value\n",
    "    # audience = 'https://my-cloud-run-service.run.app/'\n",
    "    # For Cloud Run, `endpoint` is the URL (hostname + path) receiving the request\n",
    "    # endpoint = 'https://my-cloud-run-service.run.app/my/awesome/url'\n",
    "    \n",
    "    auth_req = google.auth.transport.requests.Request()\n",
    "    id_token = google.oauth2.id_token.fetch_id_token(auth_req, endpoint)\n",
    "\n",
    "    return id_token"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0561443b-05af-444f-be74-f36d3904da7b",
   "metadata": {},
   "source": [
    "#### run_cloudsql_query()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "800ee979-1aa4-4ca5-922b-9689da575589",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sqlalchemy\n",
    "from sqlalchemy import text, exc\n",
    "import pandas as pd\n",
    "import logging\n",
    "\n",
    "async def run_cloudsql_query(pool, sql: str, params = None, output_as_df: bool = True):\n",
    "    \"\"\"Executes a SQL query or statement against the Cloud SQL database pool.\n",
    "\n",
    "    Handles various SQL statements:\n",
    "    - SELECT/WITH: Returns results as a DataFrame (if output_as_df=True)\n",
    "      or ResultProxy. Supports parameters. Does not commit.\n",
    "    - EXPLAIN/EXPLAIN ANALYZE: Executes the explain, returns the query plan\n",
    "      as a formatted multi-line string. Ignores output_as_df.\n",
    "      Supports parameters. Does not commit.\n",
    "    - INSERT/UPDATE/DELETE/CREATE/ALTER etc.: Executes the statement,\n",
    "      commits the transaction, logs info, and returns the ResultProxy.\n",
    "      Supports single or bulk parameters (executemany).\n",
    "\n",
    "    Args:\n",
    "      pool: An asynchronous SQLAlchemy connection pool.\n",
    "      sql: A string containing the SQL query or statement template.\n",
    "      params: Optional.\n",
    "        - None: Execute raw SQL (Use with caution for non-SELECT/EXPLAIN).\n",
    "        - dict or tuple: Parameters for a single execution.\n",
    "        - list of dicts/tuples: Parameters for bulk execution (executemany).\n",
    "      output_as_df (bool): If True and query is SELECT/WITH, return pandas DataFrame.\n",
    "                           Ignored for EXPLAIN and non-data-returning statements.\n",
    "\n",
    "    Returns:\n",
    "      pandas.DataFrame | str | sqlalchemy.engine.Result | None:\n",
    "        - DataFrame: For SELECT/WITH if output_as_df=True.\n",
    "        - str: For EXPLAIN/EXPLAIN ANALYZE, containing the formatted query plan.\n",
    "        - ResultProxy: For non-SELECT/WITH/EXPLAIN statements, or SELECT/WITH\n",
    "                       if output_as_df=False.\n",
    "        - None: If a SQLAlchemy ProgrammingError or other specific error occurs.\n",
    "    \"\"\"\n",
    "    sql_lower_stripped = sql.strip().lower()\n",
    "    is_select_with = sql_lower_stripped.startswith(('select', 'with'))\n",
    "    is_explain = sql_lower_stripped.startswith('explain')\n",
    "    is_data_returning = is_select_with or is_explain\n",
    "    effective_output_as_df = output_as_df and is_select_with\n",
    "    is_bulk_operation = isinstance(params, list) and len(params) > 0\n",
    "\n",
    "    async with pool.connect() as conn:\n",
    "        try:\n",
    "            if params:\n",
    "                result = await conn.execute(text(sql), params)\n",
    "            else:\n",
    "                result = await conn.execute(text(sql))\n",
    "\n",
    "            if is_data_returning:\n",
    "                if is_explain:\n",
    "                    try:\n",
    "                        plan_rows = result.fetchall()\n",
    "                        query_plan = \"\\n\".join([str(row[0]) for row in plan_rows])\n",
    "                        return query_plan\n",
    "                    except Exception as e:\n",
    "                        logging.error(f\"Error fetching/formatting EXPLAIN result: {e}\")\n",
    "                        return None\n",
    "                else:  # SELECT / WITH\n",
    "                    if effective_output_as_df:\n",
    "                        try:\n",
    "                            rows = result.fetchall()\n",
    "                            column_names = result.keys()\n",
    "                            df = pd.DataFrame(rows, columns=column_names)\n",
    "                            return df\n",
    "                        except Exception as e:\n",
    "                            logging.error(f\"Error converting SELECT result to DataFrame: {e}\")\n",
    "                            return result\n",
    "                    else:\n",
    "                        return result\n",
    "            else:\n",
    "                await conn.commit()\n",
    "                operation_type = sql.strip().split()[0].upper()\n",
    "                row_count = result.rowcount\n",
    "\n",
    "                if is_bulk_operation:\n",
    "                    print(f\"Bulk {operation_type} executed for {len(params)} items. Result rowcount: {row_count}\")\n",
    "                elif operation_type in ['INSERT', 'UPDATE', 'DELETE']:\n",
    "                    print(f\"{operation_type} statement executed successfully. {row_count} row(s) affected.\")\n",
    "                else:\n",
    "                    print(f\"{operation_type} statement executed successfully. Result rowcount: {row_count}\")\n",
    "                return result\n",
    "\n",
    "        except exc.ProgrammingError as e:\n",
    "            logging.error(f\"SQL Programming Error executing query:\\nSQL: {sql[:500]}...\\nParams (sample): {str(params)[:500]}...\\nError: {e}\")\n",
    "            return None\n",
    "        except Exception as e:\n",
    "            logging.error(f\"An unexpected error occurred during query execution:\\nSQL: {sql[:500]}...\\nError: {e}\")\n",
    "            raise\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db5cf49a-a04d-429c-a5f3-2df8e068fc81",
   "metadata": {},
   "source": [
    "## Build an Agent with MCP Toolbox Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ebd328f-1f6c-4123-9bfc-5e40f1e207bc",
   "metadata": {},
   "source": [
    "### Get the Toolbox Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "100a6f20-9494-4da1-8fed-49744f5bce53",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "toolbox_url = ! gcloud run services describe toolbox --region {region} --format 'value(metadata.annotations.\"run.googleapis.com/urls\")'\n",
    "toolbox_url = json.loads(toolbox_url[0])[0]\n",
    "print(toolbox_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a48eb1b-2ea9-4fde-93f7-6ece1032b567",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Define Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff50e697-2dfb-46ba-a5bf-e64dff1ad467",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Reference: https://googleapis.github.io/genai-toolbox/getting-started/local_quickstart/\n",
    "\n",
    "from google.adk.agents import Agent\n",
    "from google.adk.runners import Runner\n",
    "from google.adk.sessions import InMemorySessionService\n",
    "from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService\n",
    "from google import genai\n",
    "from google.genai import types\n",
    "from toolbox_core import ToolboxSyncClient\n",
    "\n",
    "import asyncio\n",
    "import os\n",
    "\n",
    "# Use GOOGLE_API_KEY or Vertex AI\n",
    "# os.environ['GOOGLE_API_KEY'] = 'your-api-key'\n",
    "\n",
    "os.environ['GOOGLE_GENAI_USE_VERTEXAI'] = 'TRUE'\n",
    "os.environ['GOOGLE_CLOUD_PROJECT']      = project_id\n",
    "os.environ['GOOGLE_CLOUD_LOCATION']     = region\n",
    "\n",
    "# Get auth token to invoke Toolbox on Cloud Run\n",
    "auth_token = get_auth_token(toolbox_url)\n",
    "\n",
    "async def invoke_agent(query: str):\n",
    "  with ToolboxSyncClient(\n",
    "        toolbox_url,\n",
    "        client_headers={\"Authorization\": f\"Bearer {auth_token}\"}\n",
    "  ) as toolbox_client:\n",
    "\n",
    "      prompt = \"\"\"\n",
    "        You're a helpful financial assistant. You handle fraud detection tasks, transaction lookups,\n",
    "        account details, loan payments, suspicious activity reports, and other tasks related to \n",
    "        financial data. Always provide a customer_id when using tools to lookup information.\n",
    "      \"\"\"\n",
    "\n",
    "      app_name = 'finance_agent'\n",
    "    \n",
    "      user_id = '123'\n",
    "\n",
    "      root_agent = Agent(\n",
    "          model='gemini-2.5-flash',\n",
    "          name=app_name,\n",
    "          description='A helpful AI assistant.',\n",
    "          instruction=prompt,\n",
    "          tools=toolbox_client.load_toolset(\"finance-toolset\"),\n",
    "      )\n",
    "\n",
    "      session_service = InMemorySessionService()\n",
    "      artifacts_service = InMemoryArtifactService()\n",
    "      session = await session_service.create_session(\n",
    "          state={}, app_name=app_name, user_id=user_id\n",
    "      )\n",
    "      runner = Runner(\n",
    "          app_name=app_name,\n",
    "          agent=root_agent,\n",
    "          artifact_service=artifacts_service,\n",
    "          session_service=session_service,\n",
    "      )\n",
    "\n",
    "      content = types.Content(role='user', parts=[types.Part(text=query)])\n",
    "      events = runner.run(session_id=session.id,\n",
    "                          user_id=user_id, new_message=content)\n",
    "\n",
    "      responses = (\n",
    "        part.text\n",
    "        for event in events\n",
    "        for part in event.content.parts\n",
    "        if part.text is not None\n",
    "      )\n",
    "\n",
    "      for text in responses:\n",
    "        print(text)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4892e166-1953-4f6f-872a-a9cd6d85e787",
   "metadata": {},
   "source": [
    "### Manually Invoke the Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fae5a659-d9cb-4d2e-89fe-1f23e993261d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "prompts = [\n",
    "    \"What are the latest trasactions for customer_id 11?\",\n",
    "    \"Tell me about the recent account transfers made by customer_id 8 and flag any suspicious activity.\",\n",
    "    \"Is anything strange about the latest tranactions made by customer_id 3?\",\n",
    "]\n",
    "\n",
    "for p in prompts:\n",
    "    print(f\"User prompt: {p}\\n\")\n",
    "    print(\"Agent response:\")\n",
    "    response = await invoke_agent(p)\n",
    "    print(\"\\n\\n\")\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d926bb0f-da94-4b76-9205-38c9640f2635",
   "metadata": {},
   "source": [
    "## Inspect Agent Processing\n",
    "\n",
    "The examples above show that the Agent has access to our databases via MCP Toolbox, because the output contains data the Agent wouldn't otherwise know. But what if we want to debug an Agent's logic or understand the specific tool calls that were made? When developing locally, we can set the log level to info to see detailed, step-by-step Agent execution and tool call details. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "142bc4bd-6ffc-45ee-a649-e6069984835c",
   "metadata": {},
   "source": [
    "### Set log level to INFO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce4f46b5-35c0-4d47-9d5a-8774aedb234b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "# Configure the root logger to output messages with INFO level or above\n",
    "logging.basicConfig(level=logging.INFO, stream=sys.stdout, format='%(asctime)s[%(levelname)5s][%(name)14s] - %(message)s',  datefmt='%H:%M:%S', force=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27d15117-2083-4f1d-8a92-3fbf9095766f",
   "metadata": {},
   "source": [
    "### Examine Detailed Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c057993a-e8c6-43e4-99ba-7f2b8b30066b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "prompt = \"How many transfers has customer_id 462 made recently? Is there anything strange about the transfers?\"\n",
    "await invoke_agent(prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeab9dd3-117d-4405-8726-9b5692754f47",
   "metadata": {},
   "source": [
    "## Deploy the ADK Agent to Cloud Run\n",
    "\n",
    "Now we will deploy the ADK agent to Cloud Run. The deployment process will use a `Dockerfile` to containerize the application and configure it to connect to the Cloud SQL instance for session management. We will also enable Cloud Logging and Cloud Trace for observability.\n",
    "\n",
    "References:\n",
    "- https://github.com/alphinside/deploy-and-manage-adk-service/tree/main/weather_agent\n",
    "- https://codelabs.developers.google.com/deploy-manage-observe-adk-cloud-run#1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8285d802-aa9c-4854-a658-66674599771e",
   "metadata": {},
   "source": [
    "### Use Cloud SQL for Session Persistence\n",
    "\n",
    "The Agent Development Kit (ADK) uses a Session Service to maintain the state of conversations. By default, it uses an in-memory session service, but this is not suitable for production environments where the server may restart, and sessions would be lost. To ensure persistence, we can connect the ADK to a Cloud SQL database to act as the persistent store for the Session service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7928c34c-00c7-4546-8845-ecc858fb5a8c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get Cloud SQL connection name\n",
    "cloud_sql_psc_ip = !gcloud compute forwarding-rules describe psc-forwarding-rule-{cloud_sql_instance} \\\n",
    "    --region={region} \\\n",
    "    --project={project_id} \\\n",
    "    --format=\"value(IPAddress)\"\n",
    "cloud_sql_psc_ip = cloud_sql_psc_ip[0]\n",
    "\n",
    "# Construct the SESSION_SERVICE_URI\n",
    "session_service_uri = f\"postgresql+pg8000://postgres:{cloud_sql_password}@{cloud_sql_psc_ip}:5432/postgres\"\n",
    "\n",
    "print(f\"SESSION_SERVICE_URI: {session_service_uri}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "084fabf9-6dcf-4aed-8e92-c7dae6bf6411",
   "metadata": {},
   "source": [
    "### Create `__init__.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2427aca-d9c5-45d1-b58c-5a699db88fc3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Initialize the agent directory\n",
    "! mkdir -p deploy\n",
    "! mkdir -p deploy/finance_agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2391e43e-8ada-4a8c-bbd5-95a6f04432d2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile deploy/finance_agent/__init__.py\n",
    "# __init__.py\n",
    "\n",
    "from .agent import root_agent\n",
    "\n",
    "__all__ = [\"root_agent\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64359e5b-f758-429f-8ebf-b3310ca9ffb6",
   "metadata": {},
   "source": [
    "### Create `agent.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e32a6a7f-9c89-4940-8904-8f12679249ea",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile deploy/finance_agent/agent.py\n",
    "import os\n",
    "import google.auth\n",
    "import google.auth.transport.requests\n",
    "import google.oauth2.id_token\n",
    "from google.adk.agents import Agent\n",
    "from toolbox_core import ToolboxSyncClient\n",
    "from opentelemetry.instrumentation.aiohttp_client import (AioHttpClientInstrumentor)\n",
    "\n",
    "# --- Environment and Authentication ---\n",
    "\n",
    "# It's best practice to set these in the Cloud Run service's environment variables\n",
    "# instead of hardcoding them in the script.\n",
    "# os.environ['GOOGLE_GENAI_USE_VERTEXAI'] = 'TRUE'\n",
    "# os.environ['GOOGLE_CLOUD_PROJECT']      = 'your-project-id'\n",
    "# os.environ['GOOGLE_CLOUD_LOCATION']     = 'your-region'\n",
    "\n",
    "def get_auth_token(endpoint: str) -> str:\n",
    "    \"\"\"Fetches an OIDC token for authenticating with a Cloud Run service.\"\"\"\n",
    "    auth_req = google.auth.transport.requests.Request()\n",
    "    return google.oauth2.id_token.fetch_id_token(auth_req, endpoint)\n",
    "\n",
    "# The URL for the toolbox service should be passed as an environment variable\n",
    "# for better flexibility between different environments (dev, prod).\n",
    "TOOLBOX_URL = os.getenv(\"TOOLBOX_URL\")\n",
    "if not TOOLBOX_URL:\n",
    "    raise ValueError(\"TOOLBOX_URL environment variable not set.\")\n",
    "\n",
    "# Get an auth token to invoke the Toolbox on Cloud Run\n",
    "auth_token = get_auth_token(TOOLBOX_URL)\n",
    "\n",
    "# --- Agent Definition ---\n",
    "\n",
    "# Initialize the Toolbox client once and reuse it\n",
    "print(f\"Connecting to toolbox server at {TOOLBOX_URL}\")\n",
    "toolbox_client = ToolboxSyncClient(\n",
    "    TOOLBOX_URL,\n",
    "    client_headers={\"Authorization\": f\"Bearer {auth_token}\"}\n",
    ")\n",
    "\n",
    "# Define the agent's instructions\n",
    "prompt = \"\"\"\n",
    "You're a helpful financial assistant. You handle fraud detection tasks, transaction lookups,\n",
    "account details, loan payments, suspicious activity reports, and other tasks related to\n",
    "financial data. Always provide a customer_id when using tools to lookup information.\n",
    "\"\"\"\n",
    "\n",
    "# Define the root agent\n",
    "root_agent = Agent(\n",
    "    model='gemini-2.5-flash',\n",
    "    name='finance_agent',\n",
    "    description='A helpful AI assistant for financial tasks.',\n",
    "    instruction=prompt,\n",
    "    # Load the tools from your MCP Toolbox\n",
    "    tools=toolbox_client.load_toolset(\"finance-toolset\"),\n",
    ")\n",
    "\n",
    "# Enable the AioHttp library for Toolbox tracing\n",
    "AioHttpClientInstrumentor().instrument()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d277e6c9-3391-4809-b781-1e4fed160d03",
   "metadata": {},
   "source": [
    "### Create `server.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "338ffbc8-ba7c-44d4-805d-4877be60eadb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile deploy/server.py\n",
    "import os\n",
    "from fastapi import FastAPI\n",
    "from google.adk.cli.fast_api import get_fast_api_app\n",
    "from google.cloud import logging as google_cloud_logging\n",
    "\n",
    "# OTEL imports\n",
    "from opentelemetry import trace\n",
    "from opentelemetry.sdk.trace import TracerProvider, export\n",
    "from opentelemetry.propagate import set_global_textmap\n",
    "from opentelemetry.propagators.cloud_trace_propagator import CloudTraceFormatPropagator\n",
    "from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter\n",
    "\n",
    "# --- Server Setup ---\n",
    "\n",
    "# Initialize the Google Cloud Logging client\n",
    "logging_client = google_cloud_logging.Client()\n",
    "logger = logging_client.logger(__name__)\n",
    "\n",
    "# The directory where your agent.py and __init__.py are located\n",
    "AGENT_DIR = os.path.dirname(os.path.abspath(__file__))\n",
    "\n",
    "# Get the session service URI from environment variables.\n",
    "# This is where you'll provide the connection string for Cloud SQL.\n",
    "session_uri = os.getenv(\"SESSION_SERVICE_URI\")\n",
    "\n",
    "# Prepare arguments for the FastAPI app\n",
    "app_args = {\"agents_dir\": AGENT_DIR, \"web\": True}\n",
    "\n",
    "# Use the Cloud SQL session service if the URI is provided\n",
    "if session_uri:\n",
    "    app_args[\"session_service_uri\"] = session_uri\n",
    "    logger.log_text(f\"Using database session service.\", severity=\"INFO\")\n",
    "else:\n",
    "    logger.log_text(\n",
    "        \"SESSION_SERVICE_URI not provided. Using in-memory session service. \"\n",
    "        \"Sessions will be lost on server restart.\",\n",
    "        severity=\"WARNING\",\n",
    "    )\n",
    "    \n",
    "# Observability\n",
    "provider = TracerProvider()\n",
    "processor = export.BatchSpanProcessor(CloudTraceSpanExporter())\n",
    "provider.add_span_processor(processor)\n",
    "trace.set_tracer_provider(provider)\n",
    "print(\"Set trace provider\")\n",
    "\n",
    "# Create the FastAPI app using the ADK utility\n",
    "app: FastAPI = get_fast_api_app(**app_args)\n",
    "\n",
    "app.title = \"finance-agent\"\n",
    "app.description = \"API for interacting with the ADK Finance Agent\"\n",
    "\n",
    "# Main execution\n",
    "if __name__ == \"__main__\":\n",
    "    import uvicorn\n",
    "    \n",
    "    # Enable the AioHttp library for Toolbox tracing\n",
    "    AioHttpClientInstrumentor().instrument()\n",
    "\n",
    "    # Instrument the FastAPI app.\n",
    "    # This automatically creates parent spans for requests and propagates the trace context.\n",
    "    FastAPIInstrumentor.instrument_app(app)\n",
    "\n",
    "    uvicorn.run(app, host=\"0.0.0.0\", port=8080)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "496c4dc8-726c-4b10-aa49-924685b81a54",
   "metadata": {},
   "source": [
    "### Create `pyproject.toml`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c83f980-53ae-4e03-aef5-64742f47e60a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile deploy/pyproject.toml\n",
    "[project]\n",
    "name = \"deploy-and-manage-adk-service\"\n",
    "version = \"0.1.0\"\n",
    "description = \"Add your description here\"\n",
    "readme = \"README.md\"\n",
    "requires-python = \">=3.12\"\n",
    "dependencies = [\n",
    "    \"google-adk==1.5.0\",\n",
    "    \"locust==2.37.10\",\n",
    "    \"pg8000==1.31.2\",\n",
    "    \"python-dotenv==1.1.0\",\n",
    "    \"toolbox-core==0.2.1\",\n",
    "    \"google-genai==1.23.0\",\n",
    "    \"opentelemetry-instrumentation-fastapi==0.55b1\",\n",
    "    \"opentelemetry-api>=1.20.0\",\n",
    "    \"opentelemetry-sdk>=1.20.0\",\n",
    "    \"opentelemetry-instrumentation-fastapi>=0.40.0\",\n",
    "    \"opentelemetry-exporter-gcp-trace>=1.9.0\",\n",
    "    \"opentelemetry-propagator-gcp==1.9.0\",\n",
    "    \"opentelemetry-exporter-gcp-monitoring==1.9.0a0\",\n",
    "    \"opentelemetry-exporter-otlp\",\n",
    "    \"opentelemetry-instrumentation-requests==0.55b1\",\n",
    "    \"opentelemetry-instrumentation-aiohttp-client\"\n",
    "]\n",
    "\n",
    "[dependency-groups]\n",
    "dev = [\n",
    "    \"pytest==8.4.0\",\n",
    "    \"ruff==0.11.13\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d98391e7-99e5-48ec-b6f2-ce95f7bd5810",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create a Dockerfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ad8a019-bd41-46b1-aa45-b56e9d54c0a7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile deploy/Dockerfile\n",
    "FROM python:3.12-slim\n",
    "RUN pip install --no-cache-dir uv==0.7.13\n",
    "WORKDIR /app\n",
    "COPY . .\n",
    "RUN uv sync\n",
    "EXPOSE 8080\n",
    "CMD [\"uv\", \"run\", \"uvicorn\", \"server:app\", \"--host\", \"0.0.0.0\", \"--port\", \"8080\"]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6c0a98f-a6ac-4447-bac7-12c988c8f4cc",
   "metadata": {},
   "source": [
    "### Deploy the Agent to Cloud Run\n",
    "\n",
    "Here, we deploy the Finance Agent as a secure Cloud Run service. \n",
    "\n",
    "Key security configurations:\n",
    "- `--network={vpc}` / `--subnet={vpc}`: Deploys the service within our private VPC.\n",
    "- `--service-account=adk-service-account`: Uses a dedicated service account with least-privilege permissions.\n",
    "- `--no-allow-unauthenticated`: Enforces IAM authentication for all invocations.\n",
    "- `--ingress=internal`: Restricts network traffic to internal sources only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10869547-4e84-4c9e-b9e0-46ac86b0cf9b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!gcloud run deploy finance-agent --quiet --no-user-output-enabled \\\n",
    "    --source ./deploy/ \\\n",
    "    --network={vpc} \\\n",
    "    --subnet={vpc} \\\n",
    "    --ingress=internal \\\n",
    "    --port 8080 \\\n",
    "    --project {project_id} \\\n",
    "    --region {region} \\\n",
    "    --no-allow-unauthenticated \\\n",
    "    --update-env-vars=TOOLBOX_URL={toolbox_url},GOOGLE_GENAI_USE_VERTEXAI=True,SESSION_SERVICE_URI={session_service_uri},GOOGLE_CLOUD_PROJECT={project_id},GOOGLE_CLOUD_LOCATION={region},GCS_BUCKET_NAME={gcs_bucket_name} \\\n",
    "    --min=1 \\\n",
    "    --service-account=adk-service-account \\\n",
    "    --memory=1Gi\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "262b0b57-0faa-4715-84df-672339c2db6c",
   "metadata": {},
   "source": [
    "### Grant finance-agent Invoke on toolbox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08d551c7-0251-4164-b6a4-71f45e2ca824",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get account finance-agent service account\n",
    "service_account = !gcloud run services describe finance-agent \\\n",
    "    --region={region} \\\n",
    "    --format=\"value(spec.template.spec.serviceAccountName)\"\n",
    "service_account = service_account[0]\n",
    "service_account = f\"{service_account}@{project_id}.iam.gserviceaccount.com\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a3745d9-0a6e-4387-8d95-fce0c0c94c46",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Grant service account invoke permissions on toolbox\n",
    "!gcloud run services add-iam-policy-binding toolbox --region={region} \\\n",
    "    --member=serviceAccount:{service_account} \\\n",
    "    --role=\"roles/run.invoker\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84f0375e-f505-44ca-996f-d6abe82f410e",
   "metadata": {},
   "source": [
    "### Enable Inter-service Private Network Connectivity\n",
    "\n",
    "By default, Cloud Run endpoints are public. Since we've locked down the Cloud Run instances to only accept traffic from internal sources, we need to create DNS records to resolve Cloud Run endpoints to private IPs and enable Google Private Access in the `demo-vpc` subnet to enable secure service-to-service communication."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9bbd0d1-39c0-4a9c-9455-943cdd51793b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!gcloud dns managed-zones create private-run-app-zone \\\n",
    "    --description=\"Private DNS zone for internal Cloud Run traffic\" \\\n",
    "    --dns-name=\"run.app.\" \\\n",
    "    --visibility=private \\\n",
    "    --networks={vpc}\n",
    "\n",
    "# Add the 'A' record\n",
    "!gcloud dns record-sets create run.app. --zone=private-run-app-zone \\\n",
    "    --type=A --ttl=300 \\\n",
    "    --rrdatas=\"199.36.153.4,199.36.153.5,199.36.153.6,199.36.153.7\"\n",
    "\n",
    "# Add the CNAME record for wildcards\n",
    "!gcloud dns record-sets create \"*.run.app.\" --zone=private-run-app-zone \\\n",
    "    --type=CNAME --ttl=300 \\\n",
    "    --rrdatas=\"run.app.\"\n",
    "\n",
    "# Enable Private Google Access on the subnet\n",
    "!gcloud compute networks subnets update {vpc} \\\n",
    "    --region={region} \\\n",
    "    --enable-private-ip-google-access"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5357ffa7-ab96-4f71-bec8-e017fb6cd017",
   "metadata": {},
   "source": [
    "## Invoke the Cloud Run Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37bbaa23-be48-4d28-a2e9-ab1c7da2c298",
   "metadata": {},
   "source": [
    "### Define Functions for Agent Invocation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a99db656-47cc-44e4-a75e-08e00013cd91",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "import json\n",
    "import uuid\n",
    "import google.auth\n",
    "import google.auth.transport.requests\n",
    "import google.oauth2.id_token\n",
    "\n",
    "# --- Session Tracking ---\n",
    "# Use a set to keep track of sessions created during this notebook's execution.\n",
    "created_sessions = set()\n",
    "\n",
    "def create_agent_session(agent_url: str, auth_token: str, app_name: str, user_id: str, session_id: str) -> bool:\n",
    "    \"\"\"\n",
    "    Explicitly creates a new session for the agent.\n",
    "    Returns True if successful, False otherwise.\n",
    "    \"\"\"\n",
    "    session_url = f\"{agent_url}/apps/{app_name}/users/{user_id}/sessions/{session_id}\"\n",
    "    headers = {\n",
    "        \"Authorization\": f\"Bearer {auth_token}\",\n",
    "        \"Content-Type\": \"application/json\"\n",
    "    }\n",
    "    payload = {\"state\": {}}\n",
    "\n",
    "    try:\n",
    "        response = requests.post(session_url, headers=headers, json=payload, timeout=60)\n",
    "        response.raise_for_status()\n",
    "        print(f\"Successfully created session: {session_id}\")\n",
    "        return True\n",
    "    except requests.exceptions.RequestException as e:\n",
    "        print(f\"\\nError creating session: {e}\")\n",
    "        if e.response is not None:\n",
    "            print(f\"Response Status: {e.response.status_code}\")\n",
    "            print(f\"Response Text: {e.response.text}\")\n",
    "        return False\n",
    "\n",
    "\n",
    "def invoke_agent_streaming(query: str, session_id: str = None, user_id: str = 'default-user') -> None:\n",
    "    \"\"\"\n",
    "    Invokes the agent on Cloud Run, creating a session only if it's new.\n",
    "    \"\"\"\n",
    "    if session_id is None:\n",
    "        session_id = str(uuid.uuid4())\n",
    "    \n",
    "    app_name = \"finance_agent\"\n",
    "\n",
    "    try:\n",
    "        agent_url = !gcloud run services describe finance-agent --region {region} --format 'value(status.url)'\n",
    "        agent_url = agent_url[0]\n",
    "        agent_auth_token = get_auth_token(agent_url)\n",
    "\n",
    "        # --- Step 1: Create Session ONLY IF IT'S NEW ---\n",
    "        if session_id not in created_sessions:\n",
    "            if create_agent_session(agent_url, agent_auth_token, app_name, user_id, session_id):\n",
    "                # Add the new session_id to our set of tracked sessions\n",
    "                created_sessions.add(session_id)\n",
    "            else:\n",
    "                return # Stop if session creation fails\n",
    "\n",
    "        # --- Step 2: Send the Message ---\n",
    "        invoke_url = f\"{agent_url}/run_sse\"\n",
    "        headers = {\n",
    "            \"Accept\": \"text/event-stream\",\n",
    "            \"Authorization\": f\"Bearer {agent_auth_token}\",\n",
    "            \"Content-Type\": \"application/json\"\n",
    "        }\n",
    "        payload = {\n",
    "            \"app_name\": app_name,\n",
    "            \"session_id\": session_id,\n",
    "            \"user_id\": user_id,\n",
    "            \"new_message\": {\n",
    "                \"role\": \"user\",\n",
    "                \"parts\": [{\"text\": query}]\n",
    "            }\n",
    "        }\n",
    "\n",
    "        with requests.post(invoke_url, headers=headers, json=payload, stream=True, timeout=300) as response:\n",
    "            response.raise_for_status()\n",
    "            print(f\"\\nUser Query: {query}\")\n",
    "            print(\"Agent Response:\")\n",
    "            for line in response.iter_lines():\n",
    "                if line and line.decode('utf-8').startswith('data: '):\n",
    "                    try:\n",
    "                        json_data = json.loads(line.decode('utf-8')[6:])\n",
    "                        text = json_data.get('content', {}).get('parts', [{}])[0].get('text', '')\n",
    "                        print(text, end='', flush=True)\n",
    "                    except (json.JSONDecodeError, IndexError):\n",
    "                        continue\n",
    "            print(\"\\n\" + \"=\"*50 + \"\\n\")\n",
    "\n",
    "    except requests.exceptions.RequestException as e:\n",
    "        print(f\"\\nError invoking agent: {e}\")\n",
    "        if e.response is not None:\n",
    "            print(f\"Response Status: {e.response.status_code}\")\n",
    "            print(f\"Response Text: {e.response.text}\")\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46e8b6ed-b61b-4cb8-94b9-b1d28ec7829e",
   "metadata": {},
   "source": [
    "### Query the Agent\n",
    "\n",
    "We'll iterate over a few three-turn conversations with different customer ID's to show the agent working and to generate some traffic to be sampled by Cloud Trace. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "328c16d3-e5de-436d-8fef-23b93de66d13",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "customer_array = [462, 1, 32]\n",
    "\n",
    "for c in customer_array:\n",
    "\n",
    "    # A unique ID for the entire conversation\n",
    "    conversation_session_id = f\"session_{uuid.uuid4()}\"\n",
    "    conversation_user_id = f\"user_{uuid.uuid4()}\"\n",
    "\n",
    "    # The first call will see the session_id is new, create it, and add it to the 'created_sessions' set.\n",
    "    invoke_agent_streaming(\n",
    "        f\"Tell me about recent transfers and transactions by customer_id {c}.\",\n",
    "        session_id=conversation_session_id,\n",
    "        user_id=conversation_user_id\n",
    "    )\n",
    "\n",
    "    # The second and subsequent calls will find the session_id in the set and skip the creation step.\n",
    "    invoke_agent_streaming(\n",
    "        \"I'm a fraud analyst investigating these transactions and transfers. Does anything suspicious stand out about them?\",\n",
    "        session_id=conversation_session_id,\n",
    "        user_id=conversation_user_id\n",
    "    )\n",
    "\n",
    "    invoke_agent_streaming(\n",
    "        \"What are some potential next steps for the investigation?\",\n",
    "        session_id=conversation_session_id,\n",
    "        user_id=conversation_user_id\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecd029be-4bf1-4ad6-b57d-d1e4f13f6609",
   "metadata": {},
   "source": [
    "## View Persisted ADK Data in Cloud SQL"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74568df4-cad1-4f3c-9915-e70642f46b77",
   "metadata": {},
   "source": [
    "### Connect to Cloud SQL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62c8c092-21c8-4c38-8ec2-18dc56bd4040",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import asyncpg\n",
    "from google.cloud.sql.connector import Connector as CloudSqlConnector\n",
    "import sqlalchemy\n",
    "from sqlalchemy.ext.asyncio import create_async_engine\n",
    "\n",
    "# Cloud SQL instance connection string\n",
    "instance_connection_string = f\"{project_id}:{region}:{cloud_sql_instance}\"\n",
    "\n",
    "\n",
    "async def init_connection_pool(connector: CloudSqlConnector, pool_size: int = 5) -> sqlalchemy.ext.asyncio.AsyncEngine:\n",
    "    \"\"\"\n",
    "    Initializes an asynchronous SQLAlchemy connection pool for a Cloud SQL instance.\n",
    "    \"\"\"\n",
    "    async def getconn() -> asyncpg.Connection:\n",
    "        conn: asyncpg.Connection = await connector.connect_async(\n",
    "            instance_connection_string,\n",
    "            \"asyncpg\",\n",
    "            user=cloud_sql_user,\n",
    "            password=cloud_sql_password,\n",
    "            db=cloud_sql_database,\n",
    "            ip_type=\"psc\"\n",
    "        )\n",
    "        return conn\n",
    "\n",
    "    pool = create_async_engine(\n",
    "        \"postgresql+asyncpg://\",\n",
    "        async_creator=getconn,\n",
    "        pool_size=pool_size,\n",
    "        max_overflow=2,\n",
    "        pool_timeout=30,\n",
    "        pool_recycle=1800,\n",
    "    )\n",
    "    return pool\n",
    "\n",
    "\n",
    "# Initialize the Cloud SQL Connector\n",
    "connector = CloudSqlConnector()\n",
    "\n",
    "# Initialize the connection pool\n",
    "cloud_sql_pool = await init_connection_pool(connector)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "549913f6-e51c-4fe6-8cc7-f67cc20d766f",
   "metadata": {},
   "source": [
    "### View Session Data\n",
    "\n",
    "When a user starts interacting with your agent, the SessionService creates a Session object (google.adk.sessions.Session). This object acts as the container holding everything related to that one specific chat thread. Read more about this data [in the docs](https://google.github.io/adk-docs/sessions/session/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d20a69a-e453-4a5e-a1c2-b561aacdd3df",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sql = \"SELECT * FROM sessions LIMIT 10;\"\n",
    "await run_cloudsql_query(cloud_sql_pool, sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35b16d01-fb3c-4764-8996-b099f455d4b7",
   "metadata": {},
   "source": [
    "### View Event Data\n",
    "\n",
    "Events are the fundamental units of information flow within the Agent Development Kit (ADK). They represent every significant occurrence during an agent's interaction lifecycle, from initial user input to the final response and all the steps in between. Read more about this data [in the docs](https://google.github.io/adk-docs/events/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e6e71b8-8763-46f6-9fb9-3da714a9ff10",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sql = \"SELECT * FROM events LIMIT 10;\"\n",
    "await run_cloudsql_query(cloud_sql_pool, sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "826d5b75-2957-4c78-bc1d-cdae5740fd31",
   "metadata": {},
   "source": [
    "### State Data\n",
    "\n",
    "Within each Session (our conversation thread), the state attribute acts like the agent's dedicated scratchpad for that specific interaction. While session.events holds the full history, session.state is where the agent stores and updates dynamic details needed during the conversation. Read more about this data [in the docs](https://google.github.io/adk-docs/sessions/state/)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04ab1968-f235-425e-b23e-2a3d9f91f38d",
   "metadata": {},
   "source": [
    "#### View User State Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b31057e-6d83-436b-959c-c15fc22226dc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sql = \"SELECT * FROM user_states LIMIT 10;\"\n",
    "await run_cloudsql_query(cloud_sql_pool, sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cca432c4-95f8-4a96-9e82-db1d8cce8a2b",
   "metadata": {},
   "source": [
    "#### View App State Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8f2d9d0-2ad0-4af7-9408-42ee9d45c001",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sql = \"SELECT * FROM app_states LIMIT 10;\"\n",
    "await run_cloudsql_query(cloud_sql_pool, sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62bfa7b5-38b2-49e1-b478-a40685997097",
   "metadata": {},
   "source": [
    "## View Telemetry"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeebb316-88fa-4a02-a1c0-4c0fedabf1bf",
   "metadata": {},
   "source": [
    "ADK supports built-in tracing. Navigate to the Cloud Trace to view traces for the Agent invocations above:\n",
    "\n",
    "Select one of the traces, and you will see an end-to-end visualization of the Agent invocation, including MCP Toolbox tool calls as shown in the example below.\n",
    "\n",
    "![Cloud Trace](img/cloud_trace.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f08e5567",
   "metadata": {},
   "source": [
    "Congratulations, you have completed Module 3! Proceed to [`4_configure_alloydb_natural_language.ipynb`](./4_configure_alloydb_natural_language.ipynb) to configure the powerful and flexible Natural Language feature of AlloyDB AI."
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-base-py",
   "name": "workbench-notebooks.m130",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m130"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "conda-base-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
